import torch.nn as nn
import json
from collections import defaultdict, OrderedDict
import hashlib

class Module:
    module_map = defaultdict(int)
    src_modules_map = defaultdict(int)
    

    def __init__(self, nn_module: nn.Module, src=False):
        # self.cls = nn_module.__class__
        property = {k: str(v) for k, v in nn_module.__dict__.items() if k not in ['_parameters','_modules', '_buffers'] and v}
        
        self.property = {
            'class': str(nn_module.__class__).split("'")[1],   # class attr
            'type': str(type(nn_module)).split("'")[1],         # ob_type
            **property
        }

        # {k: v for k, v in nn_module.__dict__.items() if not k.startswith('_')}
        self.level = 0
        self.hash = int(hashlib.md5(str(self).encode(encoding='UTF-8')).hexdigest(), 16)
        self.children = {}
        for idx, m in enumerate(nn_module.children()):
            self.children[idx] = Module(m)
            self.hash += self.children[idx].hash << idx
            self.level = max(self.level, self.children[idx].level + 1)

        self.name = self.property['class'].split('.')[-1] + str(self.hash)[-5:]
        Module.module_map[self] += 1
        if src:
            Module.src_modules_map[self] += 1

    def print(self, indent=0):
        print('\t' * indent, self.property)
        for k, v in self.children.items():
            v.print(indent + 1)

    def __hash__(self) -> int:
        return self.hash

    def __eq__(self, othr):
        return (isinstance(othr, type(self))
                and (self.property, self.children) ==
                    (othr.property, othr.children))

    def __str__(self):
        return json.dumps(self.property) 

    def draw(self):
        node_set = set()
        node_set.add(self.name)
        edge_set = set()
        for m in self.children.values():
            ns, es = m.draw()
            node_set |= ns
            edge_set |= es
            edge_set.add((self.name, m.name))
        return node_set, edge_set

    @classmethod
    def statistic(cls, *models):
        from collections import OrderedDict
        import matplotlib.pyplot as plt
        import os
        for x in models:
            cls(x)
        s = []
        s.append(f'module_count: {sum(cls.module_map.values())}')
        s.append(f'unique_count: {len(cls.module_map)}')

        module_dic = Module.sort_map(cls.module_map)
        for k, v in module_dic.items():
            s.append(f'{k.name} (lv{k.level}) cnt:{str(v)} {str(k)}')

        fig, ax = plt.subplots(figsize=(40, 10))
        plt.xticks(rotation=270)
        count = 100
        os.makedirs('output', exist_ok=True)
        
        ax.bar([x.name for x in list(module_dic.keys())[:count]], list(module_dic.values())[:count])
        plt.savefig('output/mydraw.png')

        edge_set = set()
        for m in cls.src_modules_map:
            _, es = m.draw()
            edge_set |= es

        s.append('==========draw_graphgvis==========')
        s.append('digraph G {')
        for e in edge_set:
            s.append(f'  {e[0]} -> {e[1]}')
        for m in cls.src_modules_map:
            s.append(f'  {m.name} [shape=Msquare];')
        s.append('}')
        with open('output/log', 'w') as f:
            f.write('\n'.join(s))

    @staticmethod
    def sort_map(m):
        lis = sorted(m.items(), key=lambda x: x[1], reverse=True)
        dic = OrderedDict()
        for (k, v) in lis:
            dic[k] = v
        return dic


def test():
    from torchvision.models import resnet152
    model = Module(resnet152(True))
    model.print()
    print('hashcode', model.hash)
    print('module_count', sum(Module.module_map.values()))
    print('unique_count', len(Module.module_map))

    module_list = sorted(Module.module_map.items(), key=lambda x: x[1], reverse=True)
    for (k, v) in module_list:
        print(k.name, v, k)

    print('==========draw_cs==========')
    ns, es = model.draw()
    print('\n'.join(ns))
    for e in es:
        print(e[0], e[1])

    
    print('==========draw_graphgvis==========')
    print('digraph G {')
    for e in es:
        print(f'  {e[0]} -> {e[1]}')
    print(f'  {model.name} [shape=Msquare];')
    print('}')


def draw():
    from torchvision.models import resnet18, resnet34, resnet101, resnet152, googlenet, \
                vgg11_bn, vgg13_bn, vgg16, vgg19, squeezenet1_0, squeezenet1_1
    models = [
        resnet18(),
        resnet34(),
        resnet101(),
        resnet152(),
        googlenet(),
        vgg11_bn(),
        vgg13_bn(),
        vgg16(),
        vgg19(),
        squeezenet1_0(),
        squeezenet1_1()
    ]
    Module.statistic(*models)


if __name__ == "__main__":
    # test()
    draw()